from wsgiref.util import shift_path_info
#from matplotlib.lines import _LineStyle
from netCDF4 import Dataset
import matplotlib.pyplot as plt
from matplotlib.cm import get_cmap
import cartopy.crs as crs
#from cartopy.feature import NaturalEarthFeature
import cartopy.feature as cfeature
import cartopy.io.shapereader as shapereader

from wrf import to_np, getvar, smooth2d, get_cartopy, cartopy_xlim, \
                cartopy_ylim, latlon_coords


# Open the NetCDF file
fpath = '../../WRF_output/'
fname = 'wrfout_d01_2022-06-28_00_00_00'
ncfile = Dataset(fpath + fname)

#print(ncfile.variables)
#exit()

# Get the sea level pressure
slp = getvar(ncfile, "slp")

#print(slp)

# Smooth the sea level pressure since it tends to be noisy near the
# mountains
smooth_slp = smooth2d(slp, 3, cenweight=4)

# Get the latitude and longitude points
lats, lons = latlon_coords(slp)

# Get the cartopy mapping object
cart_proj = get_cartopy(slp)

# Create a figure
fig = plt.figure(figsize=(12,6))
# Set the GeoAxes to the projection used by WRF
ax = plt.axes(projection=cart_proj)

# Make the contour outlines and filled contours for the smoothed sea level
# pressure.
ax.contour(to_np(lons), to_np(lats), to_np(smooth_slp), 10, colors="black",
            transform=crs.PlateCarree())
cntr = ax.contourf(to_np(lons), to_np(lats), to_np(smooth_slp), 10,
             transform=crs.PlateCarree(),
             cmap=get_cmap("jet"))

# add coastlines, state boarders, etc.
#states = cfeature.NaturalEarthFeature(category="cultural", scale="50m",
#                             facecolor="none",
#                             name="admin_1_states_provinces_shp")
ax.add_geometries(shapereader.Reader('./ne_50m_admin_1_states_provinces/ne_50m_admin_1_states_provinces.shp').geometries(), \
                  crs.PlateCarree(),facecolor = 'none',ec="gray",ls=":")
ax.add_feature(cfeature.BORDERS, linestyle='-')
ax.coastlines('50m', linewidth=0.8)

# Add a color bar
fig.colorbar(cntr, ax=ax, shrink=.98)

# Set the map bounds
ax.set_xlim(cartopy_xlim(smooth_slp))
ax.set_ylim(cartopy_ylim(smooth_slp))

# Add the gridlines
ax.gridlines(color="black", linestyle="dotted")

ax.set_title("Sea Level Pressure (hPa)", fontsize = 18)
#plt.title("Sea Level Pressure (hPa)")

plt.show()
